import threading
import uuid
from typing import Dict, Optional, Any, Callable, TypeVar, Generic
from enum import Enum
from datetime import datetime
import logging
from dataclasses import dataclass, field
from concurrent.futures import ThreadPoolExecutor
import traceback
import signal
import json

from calc_manager.res_helper import ResourceHelper

logger = logging.getLogger(__name__)

class TaskStatus(Enum):
    PENDING = "PENDING"
    RUNNING = "RUNNING"
    COMPLETED = "COMPLETED"
    FAILED = "FAILED"
    CANCELLED = "CANCELLED"

T = TypeVar('T')

@dataclass
class TaskInfo(Generic[T]):
    task_id: str
    status: TaskStatus
    start_time: datetime
    end_time: Optional[datetime] = None
    result: Optional[T] = None
    error: Optional[str] = None
    cancel_event: threading.Event = field(default_factory=threading.Event)

class TaskManager:
    def __init__(self, res_uri, max_workers: int = 10):
        self._tasks: Dict[str, TaskInfo] = {}
        self._lock = threading.Lock()
        self._executor = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="task_worker")
        self.res_helper = ResourceHelper(res_uri)
        
    def create_task(self, task_func: Callable[[], T], algo_id) -> str:
        """
        Create and start a new task
        
        Args:
            task_func: The function to execute in the task
            algo_id: algo id
            
        Returns:
            str: The task ID
        """
        task_id = str(uuid.uuid4())
        task_info: TaskInfo[T] = TaskInfo(
            task_id=task_id,
            status=TaskStatus.PENDING,
            start_time=datetime.now()
        )

        # async task to res
        self.res_helper.task_async({
            "tid": task_id,
            "name": f"algo {algo_id} task",
            "status": TaskStatus.PENDING.value,
            "algoId": algo_id,
            "createTime": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        })

        with self._lock:
            self._tasks[task_id] = task_info

        def wrapped_task():
            try:
                self.update_task_status(task_id, TaskStatus.RUNNING)
                result = task_func()
                self.update_task_status(task_id, TaskStatus.COMPLETED, result=result)
            except Exception as e:
                logger.error(f"Task {task_id} failed: {str(e)}\n{traceback.format_exc()}")
                self.update_task_status(task_id, TaskStatus.FAILED, error=str(e))

        self._executor.submit(wrapped_task)
        
        return task_id

    def get_task_info(self, task_id: str) -> Optional[TaskInfo]:
        """Get task information by task ID"""
        return self._tasks.get(task_id)

    def update_task_status(self, task_id: str, status: TaskStatus, 
                          result: Any = None, error: Optional[str] = None):
        """Update task status and result"""
        with self._lock:
            if task_id in self._tasks:
                task = self._tasks[task_id]
                task.status = status
                async_task = {
                    "tid": task_id,
                    "status": status.value
                }
                if status == TaskStatus.RUNNING:
                    async_task['startTime'] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                if status in [TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED]:
                    task.end_time = datetime.now()
                    async_task['endTime'] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                if result is not None:
                    task.result = result
                    async_task['result'] = json.dumps(result)
                if error is not None:
                    task.error = error
                self.res_helper.task_status_async(async_task)

    def cancel_task(self, task_id: str) -> bool:
        """Cancel a running task"""
        with self._lock:
            if task_id not in self._tasks:
                return False
            
            task = self._tasks[task_id]
            if task.status != TaskStatus.RUNNING:
                return False

            # Signal the task to cancel
            task.cancel_event.set()
            
            task.status = TaskStatus.CANCELLED
            task.end_time = datetime.now()
            return True

    def delete_task(self, task_id: str) -> bool:
        """Delete a task by its ID"""
        with self._lock:
            if task_id not in self._tasks:
                return False
            
            # If task is running, try to cancel it first
            task = self._tasks[task_id]
            if task.status == TaskStatus.RUNNING:
                self.cancel_task(task_id)
            
            # Delete the task
            del self._tasks[task_id]
            return True

    def cleanup_old_tasks(self, max_age_hours: int = 24):
        """Clean up tasks older than max_age_hours"""
        current_time = datetime.now()
        with self._lock:
            tasks_to_remove = []
            for task_id, task in self._tasks.items():
                if task.end_time:
                    age = (current_time - task.end_time).total_seconds() / 3600
                    if age > max_age_hours:
                        tasks_to_remove.append(task_id)
            
            for task_id in tasks_to_remove:
                del self._tasks[task_id]

    def shutdown(self):
        """Shutdown the task manager and cancel all running tasks"""
        # Cancel all running tasks
        with self._lock:
            for task_id, task in self._tasks.items():
                if task.status == TaskStatus.RUNNING:
                    self.cancel_task(task_id)
        
        # Shutdown thread pool
        self._executor.shutdown(wait=True, cancel_futures=True)